# -*- coding: utf-8 -*-

# __date:       2021/6/16
# __author:     YangChao
# __function:

import numpy as np

from core.log import iqlog
from core.tool.plot_data import pic_plot
from core.experiment.top_calibration import TopCalibration


class SearchCalibration(TopCalibration):

    def run(self):
        iqlog.info(f"--- LO:{self.f_lo} RF:{self.sweep_rf[0]}-{self.sweep_rf[-1]} {self.__class__.__name__} start ---")
        for f_rf in self.rf_list:
            print(f'----------- RF = {f_rf} IQ search calibration start ... -----------')
            self._init_inst(f_rf)
            param = [f_rf, self.f_lo]
            self._leak_opt()
            param_dac = self._leak_opt(False)
            param_dc = self._leak_calibration(self.scope_list[3], self.step_list[3],
                                              self.scope_list[3], self.step_list[3], True)
            param.extend(param_dac)
            param.extend(param_dc)
            self.result.append(param)
            print(f'------------ RF = {f_rf} IQ search calibration end ... ------------\n')
        self._save_result()

    def _leak_opt(self, choice: bool = True):
        x, y = self._get_param(choice)
        init_leak = self._set_param(x, y, choice)
        if choice:
            arr1, arr2 = self.dc_process, self.dc_process
            anno = 'local'
        else:
            arr1, arr2 = self.amp_process, self.phase_process
            anno = 'sideband'
        x_list, y_list, leak_list = self._search(x, y, arr1, arr2, init_leak, choice)
        data = np.array((x_list, y_list, leak_list))
        self._save_data(data, f"{anno}-search")
        pic_plot(leak_list, f'{anno} leak search RF={self.f_rf}', self.data_loc, len(leak_list) - 1)
        self._set_param(x, y, choice)
        return [x_list[-1], y_list[-1]]

    def _search(self, x, y, arr1, arr2, init_value, choice):
        if choice:
            x_min = self.dc_min
            x_max = self.dc_max
            y_min, y_max = x_min, x_max
            anno = 'dc-dc'
        else:
            x_min = self.amp_min
            x_max = self.amp_max
            y_min = self.phase_min
            y_max = self.phase_max
            anno = 'amp-phase'
        print(f'{anno}: (x, y) = {(x, y)} cur_leak = {init_value}')
        x_list = [x]
        y_list = [y]
        leak_list = [init_value]
        length = len(arr1)
        index = 0
        direction = True
        flag = True
        best_res = init_value
        while True:
            if index > length - 1:
                break
            if flag:
                if direction:
                    _x = x + arr1[index]
                else:
                    _x = x - arr1[index]
                if _x < x_min or _x > x_max:
                    if direction:
                        direction = False
                    else:
                        index += 1
                        flag = False
                        direction = True
                    continue
                temp = self._set_param(_x, y, flag)
                if temp > best_res:
                    x = _x
                    best_res = temp
                    x_list.append(x)
                    y_list.append(y)
                    leak_list.append(best_res)
                    print(f'{anno}: (x, y) = {(x, y)} cur_leak = {best_res}')
                    continue
                if direction:
                    direction = False
                else:
                    flag = False
                    direction = True
            else:
                if direction:
                    _y = y + arr2[index]
                else:
                    _y = y - arr2[index]
                if _y < y_min or _y > y_max:
                    if direction:
                        direction = False
                    else:
                        index += 1
                        flag = True
                        direction = True
                    continue
                temp = self._set_param(x, _y, flag)
                if temp > best_res:
                    y = _y
                    best_res = temp
                    x_list.append(x)
                    y_list.append(y)
                    leak_list.append(best_res)
                    print(f'{anno}:(x, y) = {(x, y)} cur_leak = {best_res}')
                    continue
                if direction:
                    direction = False
                else:
                    index += 1
                    flag = True
                    direction = True
        return x_list, y_list, leak_list
